/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_common.hpp"
#include <iostream>
#include <complex>
#include <cmath>
#include <vector>
#include <limits>

// 修复1: 正确的nearly_equal实现
template <typename T>
bool nearly_equal(T a, T b, double eps = 1e-6) {
    if constexpr (std::is_floating_point_v<T>) {
        if (std::isinf(a) || std::isinf(b)) 
            return a == b;
        if (std::isnan(a) || std::isnan(b))
            return false;
        
        return std::fabs(a - b) < eps;
    } else {
        return a == b;  // 非浮点类型直接比较
    }
}

// 修复2: 复数类型特化
template <typename T>
bool nearly_equal(std::complex<T> a, std::complex<T> b, T eps = 1e-6) {
    return nearly_equal(a.real(), b.real(), eps) &&
           nearly_equal(a.imag(), b.imag(), eps);
}

// 修复3: 解决数组越界问题
TEST(SparseBLASlevel2, dot) {
    int failures = 0;
    
    // === 修复的测试用例1: float类型 ===
    {
        const int nz = 3;
        float x[] = {1.0f, 2.0f, 3.0f};
        int idx[] = {0, 2, 4};
        float y[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
        float result;
        
        OpenSpB_sdot(OPENSPBLAS_NO_CONJ, nz, x, idx, y, 1, &result, OPENSPBLAS_ZERO_BASE);
        float expected = 1.0f*1.0f + 2.0f*3.0f + 3.0f*5.0f;
        
        if (!nearly_equal(result, expected)) {
            std::cerr << "FAIL: sdot test got " << result << ", expected " << expected << std::endl;
            failures++;
        }
    }

    // === 修复的测试用例2: double类型 ===
    {
        const int nz = 2;
        double x[] = {1.5, -2.5};
        int idx[] = {1, 3};
        double y[] = {1.0, 2.0, 3.0, 4.0};
        double result;
        
        OpenSpB_ddot(OPENSPBLAS_NO_CONJ, nz, x, idx, y, 1, &result, OPENSPBLAS_ZERO_BASE);
        double expected = 1.5*2.0 + (-2.5)*4.0;
        
        if (!nearly_equal(result, expected)) {
            std::cerr << "FAIL: ddot test got " << result << ", expected " << expected << std::endl;
            failures++;
        }
    }

    // === 修复的测试用例3: complex<float>类型 ===
    {
        const int nz = 2;
        std::complex<float> x[] = {{1.0f, 2.0f}, {3.0f, -4.0f}};
        int idx[] = {0, 2};
        std::complex<float> y[] = {{1.0f, 1.0f}, {2.0f, 0.0f}, {0.0f, 1.0f}};
        std::complex<float> result;
        
        OpenSpB_cdot(OPENSPBLAS_CONJ, nz, x, idx, y, 1, &result, OPENSPBLAS_ZERO_BASE);
        std::complex<float> expected = 
            std::conj(x[0]) * y[0] + std::conj(x[1]) * y[2];
            
        if (!nearly_equal(result, expected)) {
            std::cerr << "FAIL: cdot test got (" << result.real() << "," << result.imag() 
                      << "), expected (" << expected.real() << "," << expected.imag() << ")" << std::endl;
            failures++;
        }
    }

    // === 修复的测试用例4: complex<double>类型 ===
    {
        const int nz = 2;
        std::complex<double> x[] = {{1.0, 0.0}, {0.0, -1.0}};
        int idx[] = {0, 1};
        // 修复4: 提供足够的安全缓冲区
        std::vector<std::complex<double>> y_buffer(10, {0.0, 0.0});
        y_buffer[0] = {1.0, 1.0};
        y_buffer[2] = {2.0, 0.0};
        
        std::complex<double> result;
        
        OpenSpB_zdot(OPENSPBLAS_NO_CONJ, nz, x, idx, y_buffer.data(), 2, &result, OPENSPBLAS_ZERO_BASE);
        std::complex<double> expected = x[0] * y_buffer[0] + x[1] * y_buffer[2];
        
        if (!nearly_equal(result, expected)) {
            std::cerr << "FAIL: zdot test got (" << result.real() << "," << result.imag() 
                      << "), expected (" << expected.real() << "," << expected.imag() << ")" << std::endl;
            failures++;
        }
    }

    // 测试结果汇总
    if (failures == 0) {
        std::cout << "All tests PASSED!" << std::endl;
    } else {
        std::cout << failures << " tests FAILED!" << std::endl;
    }
}
